#!/bin/bash

voxceleb1_path=/ceph/home/zhangy20/datasets/VoxCeleb/voxceleb1
voxceleb2_path=/ceph/home/zhangy20/datasets/VoxCeleb/voxceleb2
musan_path=~/datasets/musan
rirs_path=~/datasets/RIRS_NOISES
trials_path=data/trials.lst

nnet_type=ResNet34_half # ResNet34_quarter
pooling_type=ASP
loss_type=amsoftmax

vad=false
embedding_dim=512

. ./path.sh

stage=2
echo stage $stage

# format data dir structure by soft link
if [ $stage -eq 0 ];then
	if [ ! -d data/wav_files ]; then
		mkdir -p data/wav_files
	fi

	rm -rf data/wav_files/dev
	mkdir -p data/wav_files/dev
	# format voxceleb1
	ln -s ${voxceleb1_path}/vox1_dev_wav/* data/wav_files/dev/

	# format voxceleb2
	ln -s ${voxceleb2_path}/dev/aac/* data/wav_files/dev/

	rm -rf voxceleb1_test_v2.txt
	wget https://openslr.magicdatatech.com/resources/49/voxceleb1_test_v2.txt
	mv voxceleb1_test_v2.txt data/voxceleb1_test_v2.txt
fi


# build data list
if [ $stage -eq 1 ];then
	extension=wav
	if [[ $vad == true ]];then
		echo apply vad
		extension=vad
		python3 $SPEAKER_TRAINER_ROOT/scripts/vad.py --data_dir $voxceleb1_path --num_jobs 40
		python3 $SPEAKER_TRAINER_ROOT/scripts/vad.py --data_dir $voxceleb2_path --num_jobs 40
	fi

	echo build dev data list
	python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
		--extension $extension \
		--dataset_dir data/wav_files/dev \
		--data_list_path data/dev_list.csv

	echo build musan data list
	python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
		--extension wav \
		--dataset_dir $musan_path \
		--data_list_path data/musan_list.csv

	echo build rirs data list
	python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
		--extension wav \
		--dataset_dir $rirs_path \
		--data_list_path data/rirs_list.csv

	rm -rf $trials_path
	python3 local/format_trials.py \
		--voxceleb1_root $voxceleb1_path \
		--src_trials_path data/voxceleb1_test_v2.txt \
		--dst_trials_path $trials_path
fi

# fast train
if [ $stage -eq 2 ];then
	CUDA_VISIBLE_DEVICES=3 python3 -W ignore $SPEAKER_TRAINER_ROOT/main.py \
		--nnet_type $nnet_type \
		--loss_type $loss_type \
		--pooling_type $pooling_type \
		--batch_size 200 \
		--num_workers 80 \
		--embedding_dim $embedding_dim \
		--save_top_k 50 \
		--train_list_path data/dev_list.csv \
		--musan_list_path data/musan_list.csv \
		--rirs_list_path data/rirs_list.csv \
		--max_epochs 50 \
		--max_frames 201 --min_frames 200 \
		--learning_rate 0.005 \
		--lr_step_size 2 \
		--lr_gamma 0.9 \
		--margin 0.2 \
		--distributed_backend dp \
		--trials_path $trials_path \
		--eval_interval -1 \
		--nPerSpeaker 1 \
		--reload_dataloaders_every_epoch \
		--gpus 1
fi

# fine-tune
if [ $stage -eq 3 ];then
	ckpt_path=ckpt.pt
	CUDA_VISIBLE_DEVICES=3 python3 -W ignore $SPEAKER_TRAINER_ROOT/main.py \
		--keep_loss_weight \
		--augment \
		--nnet_type $nnet_type \
		--loss_type $loss_type \
		--batch_size 100 \
		--lr_step_size 4 \
		--learning_rate 0.00005 \
		--num_workers 80 \
		--embedding_dim $embedding_dim \
		--save_top_k 20 \
		--train_list_path data/dev_list.csv \
		--musan_list_path data/musan_list.csv \
		--rirs_list_path data/rirs_list.csv \
		--checkpoint_path $ckpt_path \
		--max_epochs 20 \
		--max_frames 201 --min_frames 150 \
		--distributed_backend dp \
		--max_seg_per_spk 500 \
		--trials_path $trials_path \
		--eval_interval -1 \
		--nPerSpeaker 2 \
		--reload_dataloaders_every_epoch \
		--gpus 1
fi


if [ $stage -eq 4 ];then
	ckpt_path=/ceph/home/zhangy20/speaker_trainer/examples/VoxCeleb/verification/exp/Res2Next50_quarter_ASP_256_amsoftmax_2021-03-05-14-43-22/epoch=49_train_loss=2.12.ckpt
	CUDA_VISIBLE_DEVICES=4 python3 -W ignore $SPEAKER_TRAINER_ROOT/main.py \
		--batch_size 64 \
		--nnet_type $nnet_type \
		--pooling_type $pooling_type \
		--num_workers 100 \
		--train_list_path data/dev_list.csv \
		--trials_path data/trials.lst \
		--gpus 1 \
		--max_frames 401 --min_frames 400 \
		--checkpoint_path $ckpt_path \
		--evaluate

	rm -rf lightning_logs
fi

